/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include "../../common/data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_confusion_transpose_custom.h"
#include "tiling/platform/platform_ascendc.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void confusion_transpose_custom(GM_ADDR x, GM_ADDR y, GM_ADDR tiling);
#endif

namespace {
constexpr uint32_t USED_CORE_NUM = 1;
constexpr uint32_t TILINGDATA_SIZE = 22;
constexpr uint32_t B = 1;
constexpr uint32_t N = 2;
constexpr uint32_t S = 64;
constexpr uint32_t HNDIV = 32;
}

extern uint8_t *GenerateTiling(uint32_t b, uint32_t n, uint32_t s, uint32_t hnDiv);

int32_t main(int32_t argc, char *argv[]) {
    size_t inputSize = B * S * N * HNDIV * sizeof(uint16_t);
    size_t outputSize = B * S * N * HNDIV * sizeof(uint16_t);
    size_t tilingSize = TILINGDATA_SIZE * sizeof(uint32_t);

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputSize);
    uint8_t *y = (uint8_t *)AscendC::GmAlloc(outputSize);
    uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);

    ReadFile("../input/input_x.bin", inputSize, x, inputSize);

    uint8_t *buf = GenerateTiling(B, N, S, HNDIV);

    memcpy_s(tiling, tilingSize, buf, tilingSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(confusion_transpose_custom, USED_CORE_NUM, x, y, tiling);

    WriteFile("../output/output.bin", y, outputSize);

    free(buf);
    AscendC::GmFree((void *)x);
    AscendC::GmFree((void *)y);
    AscendC::GmFree((void *)tiling);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *xHost, *yHost;
    uint8_t *xDevice, *yDevice, *tilingDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&xHost), inputSize));
    CHECK_ACL(aclrtMallocHost((void **)(&yHost), outputSize));

    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, outputSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&tilingDevice, tilingSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("../input/input_x.bin", inputSize, xHost, inputSize);

    uint8_t *buf = GenerateTiling(B, N, S, HNDIV);

    CHECK_ACL(aclrtMemcpy(xDevice, inputSize, xHost, inputSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(tilingDevice, tilingSize, buf, tilingSize, ACL_MEMCPY_HOST_TO_DEVICE));

    // Execute the kernel
    ACLRT_LAUNCH_KERNEL(confusion_transpose_custom)
    (USED_CORE_NUM, stream, xDevice, yDevice, tilingDevice);

    // Wait for the stop event to complete
    CHECK_ACL(aclrtSynchronizeStream(stream));

    // Copy result to host memory and write to output file
    CHECK_ACL(aclrtMemcpy(yHost, outputSize, yDevice, outputSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("../output/output.bin", yHost, outputSize);

    // Clean up memory
    free(buf);

    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(tilingDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(yHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}
